//
// Created by yanhai on 2021/10/14.
//

#include <vector>
#include <cassert>
#include "base64.h"

namespace muh {
namespace encoding {

/**
 *  Truncates the byte array at index position \a pos.
 *
 *  If \a pos is beyond the end of the array, nothing happens.
 * @tparam T
 * @param v
 * @param pos
 */
template<typename T>
static void truncate(std::vector<T> &v, typename std::vector<T>::size_type pos)
{
    if (pos < v.size()) {
        v.resize(pos);
    }
}

std::string Base64::encodeToString(const std::string &src)
{
    return encodeToString(src, Base64Encoding);
}

// reference Qt5.9.9 qtbase/src/corelib/tools/qbytearray.cpp
std::string Base64::encodeToString(const std::string &src, int options)
{
#ifndef uchar
#define uchar(c)    ((unsigned char)(c))
#endif

    const char alphabet_base64[] = "ABCDEFGH" "IJKLMNOP" "QRSTUVWX" "YZabcdef"
                                   "ghijklmn" "opqrstuv" "wxyz0123" "456789+/";
    const char alphabet_base64url[] = "ABCDEFGH" "IJKLMNOP" "QRSTUVWX" "YZabcdef"
                                      "ghijklmn" "opqrstuv" "wxyz0123" "456789-_";
    const char *const alphabet = options & Base64UrlEncoding ? alphabet_base64url : alphabet_base64;
    const char padchar = '=';
    int padlen = 0;

    std::vector<char> tmp((src.length() + 2) / 3 * 4);

    int i = 0;
    char *out = tmp.data();
    while (i < src.length()) {
        // encode 3 bytes at a time
        int chunk = 0;
        chunk |= int(uchar(src.at(i++))) << 16;
        if (i == src.length()) {
            padlen = 2;
        } else {
            chunk |= int(uchar(src.at(i++))) << 8;
            if (i == src.length())
                padlen = 1;
            else
                chunk |= int(uchar(src.at(i++)));
        }

        int j = (chunk & 0x00fc0000) >> 18;
        int k = (chunk & 0x0003f000) >> 12;
        int l = (chunk & 0x00000fc0) >> 6;
        int m = (chunk & 0x0000003f);
        *out++ = alphabet[j];
        *out++ = alphabet[k];

        if (padlen > 1) {
            if ((options & OmitTrailingEquals) == 0)
                *out++ = padchar;
        } else {
            *out++ = alphabet[l];
        }
        if (padlen > 0) {
            if ((options & OmitTrailingEquals) == 0)
                *out++ = padchar;
        } else {
            *out++ = alphabet[m];
        }
    }
    assert((options & OmitTrailingEquals) || (out == tmp.size() + tmp.data()));
    if (options & OmitTrailingEquals) {
        truncate(tmp, out - tmp.data());
    }

    return std::string{tmp.begin(), tmp.end()};
}

std::vector<char> Base64::fromBase64(const std::string &base64)
{
    return fromBase64(base64, Base64Encoding);
}

// reference Qt5.9.9 qtbase/src/corelib/tools/qbytearray.cpp
std::vector<char> Base64::fromBase64(const std::string &base64, int options)
{
    unsigned int buf = 0;
    int nbits = 0;
    std::vector<char> tmp((base64.size() * 3) / 4);

    int offset = 0;
    for (int i = 0; i < base64.size(); ++i) {
        int ch = base64.at(i);
        int d;

        if (ch >= 'A' && ch <= 'Z')
            d = ch - 'A';
        else if (ch >= 'a' && ch <= 'z')
            d = ch - 'a' + 26;
        else if (ch >= '0' && ch <= '9')
            d = ch - '0' + 52;
        else if (ch == '+' && (options & Base64UrlEncoding) == 0)
            d = 62;
        else if (ch == '-' && (options & Base64UrlEncoding) != 0)
            d = 62;
        else if (ch == '/' && (options & Base64UrlEncoding) == 0)
            d = 63;
        else if (ch == '_' && (options & Base64UrlEncoding) != 0)
            d = 63;
        else
            d = -1;

        if (d != -1) {
            buf = (buf << 6) | d;
            nbits += 6;
            if (nbits >= 8) {
                nbits -= 8;
                tmp[offset++] = buf >> nbits;
                buf &= (1 << nbits) - 1;
            }
        }
    }

    truncate(tmp, offset);
    return tmp;
}

}
}
